import pickle
import pandas as pd
import math
import numpy as np
from stock import Bar
from strategy import Strategy
from backtest1 import Backtest1

def getSharpRatio(codelist, rate, startDateTime, endDateTime):
    """
      根据股票的夏普比例进行选股，用于策略
      :param codelist: list, 股票列表
      :param rate: float, 对标的无风险利率
      :param startDateTime: str, 用来计算夏普比率的初始时间
      :param endDateTime: str, 用来计算夏普比率的结束时间
      :return: DataFrame，index是股票代码，columns是夏普比率，内容是每只股票的夏普比率，按从大到小排序
      """
    sharp = {}
    for code in codelist:
        stock = Bar(code,["change_pct"])
        stock.selectStockInformation(all_data, startDateTime, endDateTime)
        day_num = stock.stock_data.shape[0]
        sharp[code] = (stock.stock_data["change_pct"].mean()*day_num-rate)/(stock.stock_data["change_pct"].std()*math.sqrt(day_num))
    sharpdf = pd.DataFrame(sharp, index=["sharp"]).T
    sharpdf = sharpdf.sort_values("sharp", ascending=False)  # 将sharpdf按夏普比率降序排序
    return sharpdf


file_path = "./data.pkl"
with open(file_path, "rb") as f:
    all_data = pickle.load(f)
codelist = all_data["2010-01-04"].index.tolist()

# 整个量化中需要用到的指标
indicatorlist = ["open", "close", "high", "low", "avg", "volume", "change_pct"]

# 获取按夏普比率排序的股票列表
stocklist = getSharpRatio(codelist, rate=0.03, startDateTime="2010-01-04", endDateTime="2010-12-31").index
# 所选择的股票数量，回测起始日期，结束日期，初始金额，选择策略编号
stock_num = 24
startDateTime = "2011-01-01"
endDateTime = "2019-12-31"
cash = 1000000
strategy = 1

# 根据夏普比率和选择的股票个数创建多个类对象
Bars = []
for i in range(stock_num):
    code = stocklist[i]
    stock = Bar(code, indicatorlist)  # 调用Bar类创建对象
    stock.selectStockInformation(all_data, startDateTime, endDateTime)  # 得到单个股票的信息的DataFrame
    Bars.append(stock)

# 假定每天无风险利率为万分之一
no_risk_rate = {}
for i in Bars[0].stock_data.index.tolist():
    no_risk_rate[i] = 0.0001

# 创建策略类的对象，根据所选策略编号调用策略，创建回测类的对象，调用回测
s = Strategy()
if strategy == 1:
    s1 = s.trade1(Bars)
    b1 = Backtest1(Bars, cash, no_risk_rate)
    df1 = b1.get_yield(s1)
elif strategy == 2:
    s2 = s.trade2(Bars)
    b2 = Backtest1(Bars, cash, no_risk_rate)
    df2 = b2.get_yield(s2)